import torch
import torch.nn as nn
import pandas as pd
import numpy as np
from datetime import datetime, timedelta

# 设置随机种子
def set_seed(seed=42):
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_seed(42)

# 自定义数据集类
class OptimizedTimeSeriesDataset:
    def __init__(self, data, seq_length, pred_length):
        self.seq_length = seq_length
        self.pred_length = pred_length
        self.samples = []
        total_length = len(data)
        load_data = data[:, 0].reshape(-1, 1)
        weather_time_data = data[:, 1:]
        for i in range(total_length - seq_length - pred_length):
            x_load = load_data[i:i + seq_length]
            x_weather_time = weather_time_data[i:i + seq_length]
            x_combined = np.concatenate([x_load, x_weather_time], axis=1)
            y = load_data[i + seq_length:i + seq_length + pred_length].flatten()
            if np.isnan(x_combined).any() or np.isnan(y).any():
                continue
            self.samples.append((x_combined, y))

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        x, y = self.samples[index]
        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)

# 增强的 LSTM 模型类
class EnhancedLSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, output_size, dropout=0.2):
        super(EnhancedLSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=False,
                            dropout=dropout if num_layers > 1 else 0)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.LeakyReLU(0.01),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, output_size)
        )
        self.init_weights()

    def init_weights(self):
        for name, param in self.named_parameters():
            if 'weight' in name:
                nn.init.xavier_uniform_(param.data)
            elif 'bias' in name:
                nn.init.constant_(param.data, 0)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        out, _ = self.lstm(x, (h0, c0))
        out = out[:, -1, :]
        if torch.isnan(out).any():
            print("!0!LSTM 输出含 NaN")
        return self.fc(out)

# 数据归一化函数
def normalize_data(df):
    df = df.copy()
    min_vals = df.min()
    max_vals = df.max()
    for col in df.columns:
        if max_vals[col] == min_vals[col]:
            df[col] = 0
        else:
            df[col] = (df[col] - min_vals[col]) / (max_vals[col] - min_vals[col])
    return df, min_vals, max_vals

# 数据反归一化函数
def denormalize_data(data, min_val, max_val):
    return data * (max_val - min_val) + min_val

# 预测函数
# 预测函数
def predict_future_load():
    # 读取数据
    df = pd.read_csv('data/Area1_Load_Weather_Time.csv')

    # 处理缺失值
    for col in ['Max_Temperature', 'Min_Temperature', 'Avg_Temperature', 'Avg_Humidity', 'Rainfall']:
        df[col] = df[col].ffill()

    df = df[df['Time'] <= '2015-01-11']
    time_column = df['Time']
    df = df.drop('Time', axis=1)

    # 数据归一化
    df_normalized, min_vals, max_vals = normalize_data(df)

    seq_length = 96 * 3  # 3天的数据作为输入
    pred_length = 96  # 每次预测24小时

    # 计算需要预测的总时间步数（7天）
    total_pred_length = 96 * 8  # 7天 * 24小时 * 4个15分钟
    num_predictions = total_pred_length // pred_length  # 需要预测的次数

    # 设备设置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载模型
    input_size = df_normalized.shape[1]
    model = EnhancedLSTMModel(input_size, 128, 2, pred_length, dropout=0.3).to(device)
    model.load_state_dict(torch.load('best_model.pt', weights_only=True))
    model.eval()

    # 初始化预测结果列表
    all_predictions = []

    # 获取最后一个输入序列
    current_sequence = df_normalized.values[-seq_length:].copy()

    # 循环预测未来7天
    for i in range(num_predictions):
        # 准备输入
        input_tensor = torch.tensor(current_sequence, dtype=torch.float32).unsqueeze(0).to(device)

        # 进行预测
        with torch.no_grad():
            prediction = model(input_tensor)
            prediction = prediction.cpu().numpy()[0]

        # 保存预测结果
        all_predictions.append(prediction)

        # 获取最后已知的天气特征值
        last_weather = current_sequence[-1, 1:].copy()

        # 创建新的预测行，使用最后已知的天气特征值
        new_prediction_row = np.zeros((pred_length, input_size))
        new_prediction_row[:, 0] = prediction  # 第一列是负荷
        new_prediction_row[:, 1:] = last_weather  # 使用最后已知的天气特征

        # 更新序列
        current_sequence = np.vstack([current_sequence[pred_length:], new_prediction_row])

    # 合并所有预测结果
    all_predictions = np.concatenate(all_predictions)

    # 反归一化
    predictions_denorm = denormalize_data(all_predictions, min_vals['Load'], max_vals['Load'])

    # 生成未来的时间序列
    last_time = pd.to_datetime(time_column.iloc[-1])
    time_index = pd.date_range(start=last_time + pd.Timedelta(minutes=15), periods=total_pred_length, freq='15min')
    time_index_str = time_index.strftime('%Y-%m-%d %H:%M:%S')

    # 保存预测结果
    prediction_df = pd.DataFrame({
        'Time': time_index_str,
        'Load': predictions_denorm
    })
    prediction_df.to_csv('data/future_7days_load_predictions.csv', index=False)
    print("预测结果已保存为 future_7days_load_predictions.csv")

if __name__ == "__main__":
    predict_future_load()